import cv2
import os
import numpy as np


def remove_problematic_images(dataset_path, output_path):
    """
    去除过于模糊、色彩过于饱和、亮度过高的图片
    :param dataset_path: 输入文件夹路径
    :param output_path: 输出文件夹路径
    """
    if not os.path.exists(output_path):
        os.makedirs(output_path)

    for filename in os.listdir(dataset_path):
        if filename.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp')):
            file_path = os.path.join(dataset_path, filename)
            image = cv2.imread(file_path)

            if image is None:
                continue

            # 检查图像模糊程度
            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
            variance = cv2.Laplacian(gray, cv2.CV_64F).var()
            if variance < 60:  # 可调节阈值
                print(f"图像模糊: {filename}")
                continue

            # 检查色彩饱和度
            hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
            saturation = hsv[:, :, 1].mean()
            if saturation > 220:  # 可调节阈值
                print(f"色彩过于饱和: {filename}")
                continue

            # 检查亮度
            brightness = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY).mean()
            if brightness > 240:  # 可调节阈值
                print(f"亮度过高: {filename}")
                continue

            # 如果图像通过所有检查，保存到输出文件夹
            output_file_path = os.path.join(output_path, filename)
            cv2.imwrite(output_file_path, image)
            # print(f"保存有效图像: {filename}")


if __name__ == "__main__":
    # dataset_path = 'path/to/your/images'  # 替换为图片文件夹路径
    # output_path = 'path/to/save/filtered_images'  # 替换为输出文件夹路径

    # 举例:
    dataset_path = './VOC2028_short/images1'  # 替换为数据集路径
    output_path = './VOC2028_short/images2'  # 替换为输出路径

    remove_problematic_images(dataset_path, output_path)